import os
from Fingerprint.util import *
import numpy as np
from obspy import read
from multiprocessing import Pool
import time



def get_haar_image(fname):
    p = params['fingerprint']
    print(p)
    min_fp_length = get_min_fp_length(params)
    # print('min_fp_length',min_fp_length)    # 12.4
    ntimes = get_ntimes(params)
    sample_haar_images = []
    st = read(fname)
    print('样例数据持续时间',st[0].stats.endtime - st[0].stats.starttime)
    # print(p['mad_sampling_rate'])  # 1

    # 不采样
    if (p['mad_sampling_rate'] == 1):
        for i in range(len(st)):
            print(i)
            if st[i].stats.endtime - st[i].stats.starttime <= min_fp_length:
                print('continue')
                continue

            # 标准化后的向量
            haar_images, nWindows, idx1, idx2, Sxx, t = feats.data_to_haar_images(st[i].data)
            sample_haar_images.append(haar_images)



    # 采样求mad
    else:
        print('采样求mad暂时未实现')

    if (len(sample_haar_images)):
        total_haar_images = np.concatenate(sample_haar_images, axis=0)
        np.save('../out/mad/mad_sample.npy', total_haar_images)
    else:
        print("WARNING: File ", fname, " NOT SAMPLED FOR MAD CALCULATION")


def get_haar_stats():
    ntimes = get_ntimes(params)
    print(ntimes)
    sample_haar_images = np.zeros([0, params['fingerprint']['nfreq'] * ntimes])
    print(sample_haar_images)
    files = ['../data/Deci5.Pick.19991015130000.CI.CDY.EHZ.sac']
    print('len(files)',len(files))
    pool = Pool(min(params['performance']['num_fp_thread'], len(files)))
    pool.map(get_haar_image, files)
    sample_haar_images = []
    for file in files:
        file_name = '../out/mad/mad_sample.npy'
        if os.path.isfile(file_name):
            partial = np.load(file_name)
            sample_haar_images.append(partial)
            os.remove(file_name)
        else:
            print ("WARNING: File not included in MAD SAMPLE", file_name)

    total_haar_images = np.concatenate(sample_haar_images, axis=0)
    return feats.compute_haar_stats(total_haar_images, type='MAD')


if __name__ == '__main__':
    t_start = time.time()
    print("hello")
    path = '../data/fp_input_CI_CDY_EHZ.json'
    params = get_params(path)
    print(params)

    feats = init_feature_extractor(params)
    print(feats.__dict__)  # 打印对象所有属性

    print(os.getcwd())  # 当前路径

    mad_folder = '../out/mad/'
    if not os.path.exists(mad_folder):
        os.makedirs(mad_folder)

    median,mad = get_haar_stats()
    print(median)
    print('=======================')
    print(mad)

    # 写入文件
    f = open('../out/mad/median_mad_74793_1024.txt','w')
    for i in range(len(median)):
        f.write('%.16f,%.16f\n'%(median[i],mad[i]))
    f.close()
    t_end = time.time()
    print('计算median和mad用时：%.2f'%(t_end-t_start))